from torch_geometric.nn import Node2Vec
from ProcessData import CodeVectorDependenceGraph
from ProcessData.PrintFormat.ColorPrint import BluePrint
from ProcessData.Process import SymbolVector
from Utils import ConfigFile
import argparse
import os
import torch
from sklearn.cluster import KMeans
from torch_geometric.transforms import NormalizeFeatures
from Metric import Metric
from Output import Out2File


def main():
    '''
    0、读取配置文件阶段
    '''
    configpath = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'config',
                              "config" + ".ini")
    parser = argparse.ArgumentParser(description="ReadConfig")
    parser.add_argument("-c", "--config", type=str, default=configpath)
    kwarg = ConfigFile(configpath).ReadConfig()

    '''
    1、数据预处理阶段
    '''


    dataset = CodeVectorDependenceGraph(kwarg["root"], transform=NormalizeFeatures())
    data = dataset[0]

    model = Node2Vec(data.edge_index, embedding_dim=128, walk_length=20,
                     context_size=10, walks_per_node=10,
                     num_negative_samples=1, p=1, q=1, sparse=True).to(kwarg['device'])

    loader = model.loader(batch_size=128, shuffle=False, num_workers=4)
    optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)

    def train():
        model.train()
        total_loss = 0
        for pos_rw, neg_rw in loader:
            optimizer.zero_grad()
            loss = model.loss(pos_rw.to(kwarg['device']), neg_rw.to(kwarg['device']))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        return total_loss / len(loader)

    @torch.no_grad()
    def test():
        model.eval()
        z = model()
        acc = model.test(z[data.train_mask], data.y[data.train_mask],
                         z[data.test_mask], data.y[data.test_mask],
                         max_iter=150)
        return acc

    for epoch in range(1, 11):
        loss = train()
    with torch.no_grad():
        z = model()

    kmeans_input = z.detach().numpy()
    kmeans = KMeans(n_clusters=14, random_state=0).fit(kmeans_input)
    preds = kmeans.predict(kmeans_input)
    '''
    3、将训练结果输出
    '''
    Out2File(out=kmeans_input, preds=preds, dataset=dataset, **kwarg)
    '''
    4、指标阶段
    '''
    BluePrint(str(Metric(kwarg["project"], kwarg["outfile_path"])))


if __name__ == "__main__":
    main()
